import json
from zhipuai import ZhipuAI

# API配置
api_key = "37c76271355e483991f6f3f8b9646ee2.X5smhyFxV0byYcJA"
base_url = "https://open.bigmodel.cn/api/paas/v4/"
model = "glm-4-flash"

# 创建ZhipuAI客户端实例
client = ZhipuAI(api_key=api_key, base_url=base_url)

def get_word_meaning(poem_content, target_word):
    """
    获取古诗词中特定词语的释义
    :param poem_content: 古诗词的完整内容
    :param target_word: 需要解释的词语
    :return: 词语的释义
    """
    prompt = f"""
    给定古诗词内容：
    {poem_content}
    请解释其中词语 "{target_word}" 的含义。
    """
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content

def translate_sentence(poem_sentence):
    """
    将古诗词句子翻译成白话文
    :param poem_sentence: 古诗词的句子
    :return: 句子的白话文译文
    """
    prompt = f"""
    请将以下古诗词句子翻译成白话文：
    {poem_sentence}
    """
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content

def infer_poem_emotion(poem_content, emotion_options):
    """
    根据古诗词内容，推理出诗人所表达的情感
    :param poem_content: 古诗词的完整内容
    :param emotion_options: 情感选项列表
    :return: 推理出的情感
    """
    emotion_list = "，".join(emotion_options)
    prompt = f"""
    给定古诗词内容：
    {poem_content}
    请根据古诗的内容，推理出诗人所表达的情感。
    情感选项：{emotion_list}
    """
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content

def process_data(input_file, output_file):
    """
    处理测试数据并生成提交文件
    :param input_file: 测试数据文件路径
    :param output_file: 提交结果文件路径
    """
    with open(input_file, 'r', encoding='utf-8') as f:
        test_data = json.load(f)

    results = []
    for idx, item in enumerate(test_data):
        result_item = {
            "idx": idx,
            "ans_qa_words": {},
            "ans_qa_sents": {},
            "choose_id": ""
        }

        # 古诗词短语内容理解
        for word in item.get("qa_words", []):
            word_meaning = get_word_meaning(item["content"], word)
            result_item["ans_qa_words"][word] = word_meaning

        # 古诗词句子理解
        for sent in item.get("qa_sents", []):
            translated_sent = translate_sentence(sent)
            result_item["ans_qa_sents"][sent] = translated_sent

        # 情感推理
        emotion_options = [option for option in item.get("choose", {}).values()]
        poem_emotion = infer_poem_emotion(item["content"], emotion_options)
        # 假设情感选项是按顺序排列的，根据返回的情感找到对应的下标
        for i, emotion in enumerate(emotion_options):
            if emotion == poem_emotion:
                result_item["choose_id"] = i
                break

        results.append(result_item)

    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=4)

if __name__ == "__main__":
    input_file = "test_data.json"  # 测试数据文件路径
    output_file = "submit.json"    # 提交结果文件路径
    process_data(input_file, output_file)
    print("处理完成，结果已保存到 {output_file}")